{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RunnableRetry\n",
    "\n",
    "- Author: [PangPangGod](https://github.com/pangpanggod)\n",
    "- Peer Review : []()\n",
    "- Proofread : [Q0211](https://github.com/Q0211)\n",
    "- This is a part of [LangChain Open Tutorial](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial)\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/13-LangChain-Expression-Language/12-RunnableRetry.ipynb)[![Open in GitHub](https://img.shields.io/badge/Open%20in%20GitHub-181717?style=flat-square&logo=github&logoColor=white)](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/13-LangChain-Expression-Language/12-RunnableRetry.ipynb)\n",
    "## Overview\n",
    "\n",
    "This tutorial covers how to use ```RunnableRetry``` to handle retry logic effectively in LangChain workflows.  \n",
    "We'll demonstrate how to configure and use ```RunnableRetry``` with examples that showcase custom retry policies to make your workflow resilient to failures.\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "- [Overview](#overview)\n",
    "- [Environement Setup](#environment-setup)\n",
    "- [What is RunnableRetry](#what-is-runnableretry)\n",
    "- [Why Use RunnableRetry](#why-use-runnableretry)\n",
    "- [Base RunnableRetry Example](#base-runnableretry-example)\n",
    "- [RunnableRetry Bind with Chains](#runnableretry-bind-with-chains)\n",
    "\n",
    "### References\n",
    "\n",
    "- [LangChain API Reference: RunnableRetry](https://python.langchain.com/api_reference/core/runnables/langchain_core.runnables.retry.RunnableRetry.html)\n",
    "- [LangChain OpenTutorial: PydanticOutputParser](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/03-OutputParser/01-PydanticOuputParser.ipynb)\n",
    "- [LangChain Document: How to structured output](https://python.langchain.com/docs/how_to/structured_output/)\n",
    "----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Set up the environment. You may refer to [Environment Setup](https://wikidocs.net/257836) for more details.\n",
    "\n",
    "**[Note]**\n",
    "- ```langchain-opentutorial``` is a package that provides a set of easy-to-use environment setup, useful functions and utilities for tutorials. \n",
    "- You can checkout the [```langchain-opentutorial```](https://github.com/LangChain-OpenTutorial/langchain-opentutorial-pypi) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install langchain-opentutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "from langchain_opentutorial import package\n",
    "\n",
    "package.install(\n",
    "    [\n",
    "        \"langchain_core\",\n",
    "        \"langchain_openai\",\n",
    "    ],\n",
    "    verbose=False,\n",
    "    upgrade=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Environment variables have been set successfully.\n"
     ]
    }
   ],
   "source": [
    "# Set environment variables\n",
    "from langchain_opentutorial import set_env\n",
    "\n",
    "set_env(\n",
    "    {\n",
    "        \"OPENAI_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_TRACING_V2\": \"true\",\n",
    "        \"LANGCHAIN_ENDPOINT\": \"https://api.smith.langchain.com\",\n",
    "        \"LANGCHAIN_PROJECT\": \"RunnableRetry\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can alternatively set API keys such as ```OPENAI_API_KEY``` in a ```.env``` file and load them.\n",
    "\n",
    "[Note] This is not necessary if you've already set the required API keys in previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load API keys from .env file\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is RunnableRetry\n",
    "\n",
    "```RunnableRetry``` is a utility provided by LangChain that allows you to add retry mechanisms for individual ```Runnable``` objects.  \n",
    "Instead of wrapping your entire workflow in retry logic, you can apply retry policies at the level of specific tasks.   \n",
    "This helps you handle transient issues, such as network errors or intermittent failures, without restarting the entire workflow.\n",
    "\n",
    "## Why use RunnableRetry\n",
    "\n",
    "By using ```RunnableRetry```, you can:\n",
    "\n",
    "- **Avoid wrapping the entire workflow with retry logic**: Instead of restarting the entire process during frequent network calls or API failures, you can retry individual ```Runnable``` units.\n",
    "- **Implement retries per task**: This enables more efficient task recovery and makes workflows more robust.\n",
    "- **Flexible implementation**: You can implement retries using ```.with_retry()``` or define a custom retry strategy by creating a ```RunnableRetry``` with specific events, such as exception types and exponential backoff."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Base RunnableRetry Example\n",
    "\n",
    "Below is a simple example to demonstrate the effectiveness of ```RunnableRetry```.  \n",
    "In this example, we simulate a task with a chance of failure and use ```RunnableRetry``` to automatically retry it up to a maximum attempts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Failed! The number is 3.\n",
      "Failed! The number is 2.\n",
      "Success! The number is 1.\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "from langchain_core.runnables import RunnableLambda\n",
    "from langchain_core.runnables.retry import RunnableRetry\n",
    "\n",
    "\n",
    "# Define a simple function with a chance of failure.\n",
    "def random_number(input):\n",
    "    number = random.randint(1, 3)\n",
    "    if number == 1:\n",
    "        print(\"Success! The number is 1.\")\n",
    "    else:\n",
    "        print(f\"Failed! The number is {number}.\")\n",
    "        raise ValueError\n",
    "\n",
    "\n",
    "# Bind the function to RunnableLambda\n",
    "runnable = RunnableLambda(random_number)\n",
    "\n",
    "# Then bind it to RunnableRetry\n",
    "runnable_with_retries = RunnableRetry(\n",
    "    bound=runnable,\n",
    "    retry_exception_types=(ValueError,),\n",
    "    max_attempt_number=5,\n",
    "    wait_exponential_jitter=True,\n",
    ")\n",
    "\n",
    "# In this example, there is no need for input, but LangChain's Runnable requires an input argument.\n",
    "# TypeError: RunnableRetry.invoke() missing 1 required positional argument: 'input'\n",
    "input = None\n",
    "runnable_with_retries.invoke(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or you can simply implemented with ```.with_retry()``` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Success! The number is 1.\n"
     ]
    }
   ],
   "source": [
    "# Bind the function to RunnableLambda\n",
    "runnable = RunnableLambda(random_number)\n",
    "\n",
    "# with .with_retry(), no need to bind with Runnableretry(bind= ...)\n",
    "runnable_with_retries = runnable.with_retry(\n",
    "    retry_if_exception_type=(ValueError,),\n",
    "    stop_after_attempt=3,\n",
    "    wait_exponential_jitter=True,\n",
    ")\n",
    "\n",
    "input = None\n",
    "runnable_with_retries.invoke(None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RunnableRetry Bind with Chains\n",
    "\n",
    "In this example, we’ll take it a step further and demonstrate how to construct a Chain using ```ChatOpenAI```. The example will show not just the basic chain setup but also how to enhance it by incorporating ```RunnableRetry``` for robust error handling and ```PydanticOutputParser``` for structured output validation.\n",
    "\n",
    "### Components Used:\n",
    "- ```RunnableRetry```: Automatically retries failed tasks to handle transient issues, such as API call failures or timeouts.\n",
    "- ```PydanticOutputParser```: Ensures the output is parsed and validated against a defined schema, making the workflow more reliable and predictable.\n",
    "\n",
    "with ```PydanticOutputParser```, check our another tutorial [here](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/03-OutputParser/01-PydanticOuputParser.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.runnables.retry import RunnableRetry\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "# Note: Each model provider may have different error classes for API-related failures.\n",
    "# For example, OpenAI uses \"InternalServerError\", while other providers may define different exceptions.\n",
    "from openai import InternalServerError\n",
    "\n",
    "# first, define model and prompt\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "prompt = PromptTemplate.from_template(\"tell me a joke about {topic}.\")\n",
    "\n",
    "# bind with RunnableRetry\n",
    "model_bind_with_retry = RunnableRetry(\n",
    "    bound=model,\n",
    "    retry_exception_types=(InternalServerError,),\n",
    "    max_attempt_number=3,\n",
    "    wait_exponential_jitter=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Why do programmers prefer dark mode?\\n\\nBecause light attracts bugs!', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 13, 'prompt_tokens': 14, 'total_tokens': 27, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_72ed7ab54c', 'finish_reason': 'stop', 'logprobs': None}, id='run-238065f6-4562-43ff-82df-81684f2f9ffb-0', usage_metadata={'input_tokens': 14, 'output_tokens': 13, 'total_tokens': 27, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = prompt | model_bind_with_retry\n",
    "chain.invoke(\"programming\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or you can use ```.with_retry()``` with Runnables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Why did the bear sit on the log?\\n\\nBecause it wanted to be a “bear-y” comfortable seat!', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 23, 'prompt_tokens': 14, 'total_tokens': 37, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_72ed7ab54c', 'finish_reason': 'stop', 'logprobs': None}, id='run-d5f7f349-a301-447b-b908-4bd8cfd075d7-0', usage_metadata={'input_tokens': 14, 'output_tokens': 23, 'total_tokens': 37, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = prompt | model.with_retry(\n",
    "    retry_if_exception_type=(InternalServerError,),\n",
    "    stop_after_attempt=3,\n",
    "    wait_exponential_jitter=True,\n",
    ")\n",
    "chain.invoke(\"bear\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Combining ```RunnableRetry``` with ```.with_structured_output()```\n",
    "\n",
    "Methods like ```.bind()``` or ```.with_retry()``` create a **new ```Runnable``` object**, making the original instance's chainable methods (e.g., ```.with_structured_output()```) unavailable.\n",
    "\n",
    "To retain structured output functionality:\n",
    "1. **Apply ```.with_structured_output()``` first** to include parsing logic in the LLM instance.\n",
    "2. Then wrap it with ```RunnableRetry``` or ```.with_retry()``` to add retry or chaining logic.\n",
    "\n",
    "### Why this order is important\n",
    "- **Loss of Methods**: The new ```Runnable``` created by ```RunnableRetry``` or ```.with_retry()``` doesn't have the original methods.\n",
    "- **Ensure Structured Output**: Adding ```.with_structured_output()``` before retry logic ensures structured parsing stays intact while allowing retries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "# Let's assume that you want result to be parsed in this format.\n",
    "class Joke(BaseModel):\n",
    "    setup: str = Field(description=\"question to set up a joke\")\n",
    "    punchline: str = Field(description=\"answer to resolve the joke\")\n",
    "\n",
    "\n",
    "prompt = PromptTemplate.from_template(\"tell me a joke about {topic}.\")\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# bind structured output\n",
    "model_with_structured_output = model.with_structured_output(Joke)\n",
    "\n",
    "model_bind_with_retry = RunnableRetry(\n",
    "    # apply .with_structured_output() first, then use RunnableRetry to add retry logic without losing this capability.\n",
    "    bound=model_with_structured_output,\n",
    "    retry_exception_types=(InternalServerError,),\n",
    "    max_attempt_number=3,\n",
    "    wait_exponential_jitter=True,\n",
    ")\n",
    "\n",
    "chain = prompt | model_bind_with_retry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why don’t penguins like talking to strangers at parties?', punchline='They find it hard to break the ice!')"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"topic\": \"penguin\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "or you can using ```with_structured_output()``` method  with ```with_retry()``` method like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Joke(BaseModel):\n",
    "    setup: str = Field(description=\"question to set up a joke\")\n",
    "    punchline: str = Field(description=\"answer to resolve the joke\")\n",
    "\n",
    "\n",
    "prompt = PromptTemplate.from_template(\"tell me a joke about {topic}.\")\n",
    "\n",
    "model = ChatOpenAI(model=\"gpt-4o-mini\")\n",
    "\n",
    "# bind structured output\n",
    "model_with_structured_output = model.with_structured_output(Joke)\n",
    "\n",
    "# apply .with_structured_output() first, then use .with_retry() to add retry logic without losing this capability.\n",
    "model_bind_with_retry = model_with_structured_output.with_retry(\n",
    "    retry_if_exception_type=(InternalServerError,),\n",
    "    stop_after_attempt=3,\n",
    "    wait_exponential_jitter=True,\n",
    ")\n",
    "\n",
    "chain = prompt | model_bind_with_retry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Joke(setup='Why did the paper clip apply for a job?', punchline='Because it wanted to hold things together!')"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain.invoke({\"topic\": \"clip\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** If the model doesn’t support ```.with_structured_output()``` or you want to use a custom parsing method,   \n",
    "refer to the [LangChain documentation on advanced structured output](https://python.langchain.com/docs/how_to/structured_output/) for more details."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
